import sys
import time
import numpy as np
import pyautogui
import torch
import win32con
import win32gui
import win32ui
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5.QtWidgets import QApplication

from designer_demo import Ui_MainWindow
from models.common import DetectMultiBackend
from utils.dataloaders import LoadImages
from utils.general import (LOGGER, check_img_size, cv2, non_max_suppression, scale_coords)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device, time_sync

# 截取手牌的 roi
roi_current_cards = [0, 900, 1920, 1080 - 900]
# 切牌状态检测 roi
roi_abandon_status = [1460, 760, 140, 140]
# 标签 classes 位置
p_classes_path = "../data/labels/classes.txt"
# 标签 labels 位置
p_labels_path = "../data/labels/labels.txt"
# 手牌数量
cards_num_limit = [13, 14]
class UI(QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        super(UI, self).__init__(parent)
        self.setupUi(self)
        self.start_paint = False
        self.draw_arr = []
        self.current_cards = [[], [], [], [], []]
        # 顺子
        self.straight_arr = [[], [], []]
        # 刻子
        self.carving_arr = [[], [], [], [], []]
        # 可能成为顺子
        self.maybe_straight_arr = [[], [], []]
        # 可能成为刻子
        self.maybe_carving_arr = [[], [], [], [], []]
        # 被挑选后剩余可打的牌
        self.surplus_arr = [[], [], [], [], []]
        desk = QApplication.desktop()
        self.setFixedSize(desk.width(), desk.height())
        self.setWindowFlags(self.windowFlags() | Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.model = SmartMa()
        Task(parent=self).start()  # 开始任务

    def paintEvent(self, event):
        if (self.start_paint):
            # print("绘图一次")
            self.draw_obj()
            self.start_paint = False

    def draw_obj(self):
        painter = QPainter()
        painter.begin(self)
        pen01 = QPen(Qt.red, 3)
        pen02 = QPen(Qt.yellow, 3)
        painter.setFont(QFont('宋体', 13, 1))
        for (x1, y1, x2, y2), label in self.draw_arr:
            w = x2 - x1
            h = y2 - y1
            dx = roi_current_cards[0] + x1
            dy = roi_current_cards[1] + y1
            print("dx{}, dy{}".format(dx, dy))
            # painter.setPen(pen01)
            # painter.drawRect(QRectF(dx, dy, w, h))
            label_cn = self.class2label(label.split(" ")[0])
            painter.setPen(pen02)
            painter.drawText(QRectF(dx, dy - 27, w, 100), label_cn)
        painter.end()

    # 标签转中文
    def class2label(self, cls):
        classes = [i.split("\n")[0] for i in open(p_classes_path, encoding="utf-8").readlines()]
        labels = [i.split("\n")[0] for i in open(p_labels_path, encoding="utf-8").readlines()]
        return labels[classes.index(cls)]

    # 将手牌进行排序
    def order_cards(self):
        [i.sort() for i in self.current_cards]

    def init_arrs(self):
        self.surplus_arr = [[], [], [], [], []]
        self.current_cards = [[], [], [], [], []]
        self.draw_arr.clear()
        self.straight_arr = [[], [], []]
        self.carving_arr = [[], [], [], [], []]
        self.maybe_straight_arr = [[], [], []]
        self.maybe_carving_arr = [[], [], [], [], []]


    # 将获取到的牌的索引变换为实际的花色和牌数
    def index2CC(self, index):
        # 0-8 万 9-17 条 18-26 dot 27-30 东南西北 31-33中发白
        character = [i for i in range(0, 9)]
        bamboo = [i for i in range(9, 18)]
        dot = [i for i in range(18, 27)]
        wind = [i for i in range(27, 31)]
        cr = [i for i in range(31, 34)]
        if index in character:
            return 0, index
        elif index in bamboo:
            return 1, index - 9
        elif index in dot:
            return 2, index - 18
        elif index in wind:
            return 3, index - 27
        elif index in cr:
            return 4, index - 31
        else:
            return ()

    # 检测
    def detect_img(self, source):
        model = self.model.model
        imgsz = [640, 640]
        conf_thres = 0.80  # 置信度阈值
        iou_thres = 0.45
        max_det = 14  # 最大检测数量
        view_img = False
        save_crop = False
        nosave = False
        classes = None
        agnostic_nms = False
        augment = False
        visualize = False
        line_thickness = 1  # 线厚度 px
        hide_labels = False  # 隐藏标签
        hide_conf = False  # 隐藏置信度
        half = False
        print(source)
        if source == "":
            print("图片不存在")
        else:
            source = str(source)
            device = select_device(self.model.device)
            stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
            imgsz = check_img_size(imgsz, s=stride)  # 图片尺寸检查
            save_img = not nosave and not source.endswith('.txt')
            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
            dt, seen = [0.0, 0.0, 0.0], 0
            print("pred", dataset)
            for path, im, im0s, vid_cap, s in dataset:
                t1 = time_sync()
                im = torch.from_numpy(im).to(device)
                im = im.half() if half else im.float()
                im /= 255
                if len(im.shape) == 3:
                    im = im[None]
                t2 = time_sync()
                dt[0] += t2 - t1
                pred = model(im, augment=augment, visualize=visualize)
                t3 = time_sync()
                dt[1] += t3 - t2
                pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
                dt[2] += time_sync() - t3
                print("pred", pred)
                for i, det in enumerate(pred):
                    seen += 1
                    p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
                    s += '%gx%g ' % im.shape[2:]
                    annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                    if len(det):
                        det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
                        for c in det[:, -1].unique():
                            n = (det[:, -1] == c).sum()
                            s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "
                        # 这里限制只有当手牌数量等于我们的指定数量时才进行授牌的获取更新
                        if len(det) not in cards_num_limit:
                            continue
                        print("识别到的牌数", len(det))
                        for *xyxy, conf, cls in reversed(det):
                            if save_img or save_crop or view_img:
                                c = int(cls)
                                label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                                print("data", xyxy, label)
                                self.draw_arr.append((xyxy, label))
                                cx = (xyxy[2] - xyxy[0]) / 2 + xyxy[0] + roi_current_cards[0]
                                cy = (xyxy[3] - xyxy[1]) / 2 + xyxy[1] + roi_current_cards[1]
                                real_card = self.index2CC(c)
                                if len(real_card) != 0:
                                    self.current_cards[real_card[0]].append([real_card[1], (cx.item(), cy.item())])
                                    # 这里构造 [牌数, (中心x, 中心y)] 作为一张牌插入到手牌队列
                                annotator.box_label(xyxy, label, color=colors(c, True))
                        self.repaint()
                    # LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
                    self.order_cards()  # 将手牌排序
                    print(self.current_cards)
                    # im0 = annotator.result()
                    # cv2.imwrite("./tmp/single_result.jpg", im0)

    # 寻找顺子
    # 代办：12223时进行决策是否提取123
    def find_straight(self):
        for idxi, i in enumerate(self.current_cards):
            if idxi >= 3:
                break
            idxj = 0
            while True:
                cards_len = len(i)
                if cards_len < 3 or idxj > cards_len - 3:
                    break
                next_card = i[idxj + 1]
                next_next_card = i[idxj + 2]
                if (next_card[0] - i[idxj][0]) == 1 and (next_next_card[0] - next_card[0]) == 1:
                    self.straight_arr[idxi].append((i[idxj], next_card, next_next_card))
                    # 若满足则添加到对应数组
                    i.pop(idxj)
                    i.pop(idxj)
                    i.pop(idxj)
                    cards_len -= 3
                    idxj -= 1
                idxj += 1
        return self.current_cards

    # 寻找刻子
    def find_carving(self):
        for idxi, i in enumerate(self.current_cards):
            idxj = 0
            while True:
                cards_len = len(i)
                if cards_len < 3 or idxj > cards_len - 3:
                    break
                next_card = i[idxj + 1]
                next_next_card = i[idxj + 2]
                if (next_card[0] - i[idxj][0]) == 0 and (next_next_card[0] - next_card[0]) == 0:
                    self.carving_arr[idxi].append((i[idxj], next_card, next_next_card))
                    # 若满足则添加到对应数组
                    i.pop(idxj)
                    i.pop(idxj)
                    i.pop(idxj)
                    cards_len -= 3
                    idxj -= 1
                idxj += 1
        return self.current_cards

    # 寻找可能成为顺子
    def find_maybe_straight(self):
        for idxi, i in enumerate(self.current_cards):
            if idxi >= 3:
                break
            idxj = 0
            while True:
                cards_len = len(i)
                if cards_len < 2 or idxj > cards_len - 2:
                    break
                next_card = i[idxj + 1]
                if (next_card[0] - i[idxj][0]) in [1, 2]:
                    self.maybe_straight_arr[idxi].append((i[idxj], next_card))
                    # 若满足则添加到对应数组
                    i.pop(idxj)
                    i.pop(idxj)
                    cards_len -= 2
                    idxj -= 1
                idxj += 1
        return self.current_cards

    # 寻找可能成为刻子
    def find_maybe_carving(self):
        for idxi, i in enumerate(self.current_cards):
            idxj = 0
            while True:
                cards_len = len(i)
                if cards_len < 2 or idxj > cards_len - 2:
                    break
                next_card = i[idxj + 1]
                if next_card[0] == i[idxj][0]:
                    self.maybe_carving_arr[idxi].append((i[idxj], next_card))
                    # 若满足则添加到对应数组
                    i.pop(idxj)
                    i.pop(idxj)
                    cards_len -= 2
                    idxj -= 1
                idxj += 1
        return self.current_cards

    # 判断牌型
    def judge_status(self):
        sarr = [len(k) for k in self.straight_arr]
        slen = sum(sarr)
        print("当前手牌的顺子数量", slen)
        carr = [len(k) for k in self.carving_arr]
        clen = sum(carr)
        print("当前手牌的刻子数量", clen)
        darr = [len(k) for k in self.maybe_carving_arr]
        dlen = sum(darr)
        print("当前手牌的对子数量", dlen)
        if (slen + clen) == 4 and dlen == 1:
            print("达成胡牌条件")

    def abandon(self):
        aarr = [len(k) for k in self.surplus_arr]
        aidx = [idx if i >= 1 else -1 for idx, i in enumerate(aarr)]
        aidx.sort(reverse=True)
        if (aidx[0] != -1):
            print("弃牌队列", self.surplus_arr, "弃牌索引：", aidx, "打出的牌：", self.surplus_arr[aidx[0]][0])

            return self.surplus_arr[aidx[0]][0]
        else:
            print("弃牌队列", self.surplus_arr, "弃牌索引：", aidx, "出错了啊铁子")

            return ()

    def go(self):
        cards = self.find_straight()
        print("找顺子后的手牌", cards, "顺子数组", self.straight_arr)
        cards = self.find_carving()
        print("找刻子后的手牌", cards, "刻子数组", self.carving_arr)
        cards = self.find_maybe_straight()
        print("找可能顺子后的手牌", cards, "可能顺子数组", self.maybe_straight_arr)
        cards = self.find_maybe_carving()
        print("找可能刻子后的手牌", cards, "可能刻子数组", self.maybe_carving_arr)
        self.surplus_arr = cards
        print("可打的手牌", self.surplus_arr)


class Task(QThread):
    def __init__(self, *args, **kwargs):
        super(Task, self).__init__()
        self.parent = kwargs.get('parent')

    def run(self):
        ai = self.parent.model
        while True:
            # t1 = time_sync()
            # cv2.imwrite("./tmp/img.png", winShot(roi_current_cards))
            # t2 = time_sync()
            # # print("截图耗时：{}".format(t2 - t1))
            # test001("./tmp/img.png", self.parent)
            # self.parent.start_paint = True
            time.sleep(0.5)
            print(ai.checkStatus())
            self.parent.init_arrs()
            if ai.checkStatus() == 1:
                # 该我们切牌
                # self.parent.start_paint = False
                # self.parent.update()
                cv2.imwrite("./tmp/img.png", winShot(roi_current_cards))
                self.parent.start_paint = True
                self.parent.detect_img("./tmp/img.png")
                # 这里我们已经得到了手牌序列以及将其排序
                self.parent.go()
                # 对已有的牌进行分析
                ac = self.parent.abandon()
                if len(ac) != 0:
                    print(ac)
                    x, y = ac[1]
                    pyautogui.moveTo(x, y)
                    pyautogui.sleep(0.5)
                    pyautogui.click(x, y)
                    pyautogui.moveTo(10, 1070)




class SmartMa():
    def __init__(self):

        # 手牌, 0万 1条 2筒 3字牌(0东 1南 2西 3北 4中 5发 6白)
        self.current_arr = [[3, 5], [2, 3, 4, 6], [1, 1, 1, 2], [5, 5, 5]]
        self.device = "cpu"
        self.model_path = "best.pt"
        self.model = self.model_load(weights=self.model_path, device=self.device)

    '''
        状态检测
        0 不做任何操作 1 该我们切牌
    '''

    def checkStatus(self):
        # 我们认为指定区域出现白色说明轮到我们出牌
        img = winShot(roi_abandon_status)
        print(img.max())
        # 大于220说明该我们切牌
        if img.max() >= 220:
            # 该我们切牌
            return 1
        return 0

    # 模型初始化
    def model_load(self, weights="", device=''):
        device = select_device(device)
        model = DetectMultiBackend(weights, device=device, dnn=False)
        stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
        if pt:
            model.model.float()
        print("模型加载完成")
        return model


# 指定窗口截图
def winShot(roi, fmt="RGB"):
    left, top, w, h = roi
    # 这里的 top 似乎是 bottom 一样的存在，
    hwnd = win32gui.FindWindow(None, "雀魂麻将")
    # hwnd = win32gui.GetDesktopWindow()
    wDC = win32gui.GetWindowDC(hwnd)
    dcObj = win32ui.CreateDCFromHandle(wDC)
    cDC = dcObj.CreateCompatibleDC()
    dataBitMap = win32ui.CreateBitmap()

    dataBitMap.CreateCompatibleBitmap(dcObj, w, h)

    cDC.SelectObject(dataBitMap)
    cDC.BitBlt((0, 0), (w, h), dcObj, (left, top), win32con.SRCCOPY)
    signedIntsArray = dataBitMap.GetBitmapBits(True)
    img = np.frombuffer(signedIntsArray, dtype="uint8")
    img.shape = (h, w, 4)

    # Free Resources
    dcObj.DeleteDC()
    cDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, wDC)
    win32gui.DeleteObject(dataBitMap.GetHandle())

    if fmt == 'BGR':
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGBA2BGR)
    if fmt == 'RGB':
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGBA2RGB)
    else:
        raise ValueError('erro!')


def test():
    sm = SmartMa()
    sm.go()
    sm.judge_status()
    sm.abandon()


def test001(img_path, model):
    t1 = time_sync()
    model.detect_img(img_path)
    t2 = time_sync()
    # print("识别耗时", t2 - t1)


def test003():
    img = pyautogui.screenshot()
    open_cv_image = np.array(img)
    open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
    # cv2.imwrite('image.jpg', open_cv_image)
    cv2.imshow("test", open_cv_image)


def test004():
    app = QApplication(sys.argv)
    ui = UI()
    ui.show()
    ui.start_paint = True
    ui.repaint()
    sys.exit(app.exec_())


def test005():
    AI = SmartMa()
    AI.checkStatus()


if __name__ == '__main__':
    test004()
